package manager

import (
	"context"
	"errors"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
	"time"
)

// WebsocketClient 每个用户的连接
type WebsocketClient struct {
	Conn   *websocket.Conn
	Ctx    context.Context
	Cancel context.CancelFunc
}

// WebsocketManager 管理连接
type WebsocketManager struct {
	ClientMap map[string]*WebsocketClient
	mu        sync.Mutex
}

func NewWebsocketManager() *WebsocketManager {
	return &WebsocketManager{
		ClientMap: make(map[string]*WebsocketClient),
	}
}

func (w *WebsocketManager) NewClient(c *gin.Context) (*WebsocketClient, error) {
	userId := w.GetUserId(c)

	var upGrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		Subprotocols: []string{userId},
	}

	conn, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(context.Background())
	client := &WebsocketClient{
		Conn:   conn,
		Ctx:    ctx,
		Cancel: cancel,
	}

	//关闭回调
	conn.SetCloseHandler(func(code int, text string) error {
		log.Printf("websocket [%v] [%#v]关闭连接 \n", c.Request.RequestURI, conn.RemoteAddr().String())
		client.Cancel()
		w.mu.Lock()
		delete(w.ClientMap, userId)
		w.mu.Unlock()

		message := websocket.FormatCloseMessage(code, "")
		err := conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(time.Second))
		if err != nil {
			log.Printf("websocket [%v] [%#v]关闭连接失败[%#v] \n", c.Request.RequestURI, conn.RemoteAddr().String(), err)
			return err
		}
		return nil
	})

	log.Printf("websocket [%v] [%#v]建立连接 \n", c.Request.RequestURI, conn.RemoteAddr().String())

	return client, nil
}

func (w *WebsocketManager) NewWebsocketConnection(c *gin.Context) (*WebsocketClient, error) {
	userId := w.GetUserId(c)

	client, err := w.NewClient(c)
	if err != nil {
		return nil, err
	}
	w.mu.Lock()
	w.ClientMap[userId] = client
	w.mu.Unlock()

	return client, err
}

func (w *WebsocketManager) GetWebsocketConnections(c *gin.Context) (*WebsocketClient, error) {
	userId := w.GetUserId(c)
	client, ok := w.ClientMap[userId]
	if client == nil || !ok {
		return nil, errors.New("websocket 没有可用连接")
	}
	return client, nil
}

func (w *WebsocketManager) GetUserId(c *gin.Context) string {
	return c.Query("userId")
}
